#!/usr/bin/env python3
"""
增强工作流池使用示例
展示如何使用增强后的工作流池功能，包括动态扩缩容、健康检查和负载均衡
"""

import asyncio
import sys
import os

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from src.research_core.workflow_pool import WorkflowPool, LoadBalancer, initialize_workflow_pool, shutdown_workflow_pool, submit_research_task


async def example_callback(result):
    """示例回调函数"""
    print(f"任务完成，结果: {result.get('final_answer', '无答案')[:100]}...")


async def main():
    """主函数演示各种工作流池功能"""
    print("=== 增强工作流池使用示例 ===\n")
    
    # 1. 初始化工作流池
    print("1. 初始化工作流池...")
    await initialize_workflow_pool()
    
    # 2. 提交不同类型的任务
    print("\n2. 提交不同类型的任务...")
    
    # 简单问题 - 应该使用单代理工作流
    await asyncio.sleep(1)  # 等待前面的任务入队
    await submit_research_task(
        "什么是人工智能？", 
        example_callback,
        {"timeout": 60}
    )
    
    # 中等复杂问题 - 应该使用标准多代理工作流
    await asyncio.sleep(1)  # 等待前面的任务入队
    await submit_research_task(
        "机器学习和深度学习有什么区别和联系？", 
        example_callback,
        {"timeout": 120}
    )
    
    # 复杂问题 - 应该使用增强多代理工作流
    await asyncio.sleep(1)  # 等待前面的任务入队
    await submit_research_task(
        "请详细分析当前大语言模型的发展趋势、技术挑战以及在各行业的应用前景，并给出相关的数据支持。", 
        example_callback,
        {"timeout": 180}
    )
    
    # 3. 使用负载均衡器
    print("\n3. 使用负载均衡器...")
    load_balancer = LoadBalancer(pool_size=2, max_workers_per_pool=3)
    await load_balancer.start_all()
    
    # 提交批量任务
    questions = [
        "云计算的优势是什么？",
        "区块链技术的应用场景有哪些？",
        "物联网的发展现状如何？",
        "5G技术对社会的影响是什么？"
    ]
    
    for question in questions:
        await load_balancer.submit(question, example_callback)
        await asyncio.sleep(0.5)  # 稍微间隔一下
    
    # 4. 获取统计信息
    print("\n4. 获取工作流池统计信息...")
    from src.research_core.workflow_pool import global_workflow_pool
    stats = global_workflow_pool.get_stats()
    print(f"工作流池状态: {'运行中' if stats['running'] else '已停止'}")
    print(f"工作线程数: {stats['workers']}")
    print(f"任务队列大小: {stats['queue_size']}")
    print(f"成功率: {stats['success_rate']:.2f}%")
    print(f"平均处理时间: {stats['average_processing_time']:.2f}秒")
    
    # 5. 等待任务完成
    print("\n5. 等待任务完成...")
    await asyncio.sleep(15)
    
    # 6. 再次获取统计信息
    print("\n6. 任务完成后的统计信息...")
    stats = global_workflow_pool.get_stats()
    print(f"已提交任务数: {stats['stats']['tasks_submitted']}")
    print(f"已完成任务数: {stats['stats']['tasks_completed']}")
    print(f"失败任务数: {stats['stats']['tasks_failed']}")
    print(f"平均处理时间: {stats['average_processing_time']:.2f}秒")
    
    # 7. 关闭负载均衡器
    print("\n7. 关闭负载均衡器...")
    await load_balancer.stop_all()
    
    # 8. 关闭工作流池
    print("\n8. 关闭工作流池...")
    await shutdown_workflow_pool()
    
    print("\n=== 示例结束 ===")


if __name__ == "__main__":
    asyncio.run(main())